package net.cyue.ort.llm.tokenizer;

import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.PropertyNamingStrategies;
import com.fasterxml.jackson.databind.annotation.JsonNaming;
import net.cyue.ort.llm.template.ModelChatTemplate;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.util.List;
import java.util.Map;


@JsonIgnoreProperties(ignoreUnknown = true)
@JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
public class TransformersTokenizerConfig {
    private boolean addBosToken;
    private String bosToken;
    private String eosToken;
    private String padToken;
    private String unkToken;
    private String chatTemplate;
    private int modelMaxLength;
    private String tokenizerClass;
    private Map<String, AddedToken> addedTokensDecoder;
    private List<String> additionalSpecialTokens;

    @JsonIgnoreProperties(ignoreUnknown = true)
    @JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
    public static class AddedToken {
        private String content;
        private boolean lstrip;
        private boolean normalized;
        private boolean rstrip;
        private boolean singleWord;
        private boolean special;

        // Getters and setters
        public void setContent(String content) { this.content = content; }
        public void setLstrip(boolean lstrip) { this.lstrip = lstrip; }
        public void setNormalized(boolean normalized) { this.normalized = normalized; }
        public void setRstrip(boolean rstrip) { this.rstrip = rstrip; }
        public void setSingleWord(boolean singleWord) { this.singleWord = singleWord; }
        public void setSpecial(boolean special) { this.special = special; }

        public String getContent() { return content; }
        public boolean isLstrip() { return lstrip; }
        public boolean isNormalized() { return normalized; }
        public boolean isRstrip() { return rstrip; }
        public boolean isSingleWord() { return singleWord; }
        public boolean isSpecial() { return special; }
    }

    private TransformersTokenizerConfig() {}
    private String extractTokenContent(Object tokenObj) {
        if (tokenObj == null) {
            return null;
        }
        if (tokenObj instanceof String) {
            return (String) tokenObj;
        }
        if (tokenObj instanceof Map) {
            @SuppressWarnings("unchecked")
            Map<String, Object> tokenMap = (Map<String, Object>) tokenObj;
            return (String) tokenMap.get("content");
        }
        return null;
    }

    // Getters and setters

    public void setAddBosToken(boolean addBosToken) { this.addBosToken = addBosToken; }
    public void setBosToken(String bosToken) { this.bosToken = bosToken; }
    public void setEosToken(String eosToken) { this.eosToken = eosToken; }
    public void setPadToken(String padToken) { this.padToken = padToken; }
    public void setUnkToken(String unkToken) { this.unkToken = unkToken; }
    public void setChatTemplate(String chatTemplate) { this.chatTemplate = chatTemplate; }
    public void setModelMaxLength(int modelMaxLength) { this.modelMaxLength = modelMaxLength; }
    public void setTokenizerClass(String tokenizerClass) { this.tokenizerClass = tokenizerClass; }
    public void setAddedTokensDecoder(Map<String, AddedToken> addedTokensDecoder) { this.addedTokensDecoder = addedTokensDecoder; }
    public void setAdditionalSpecialTokens(List<String> additionalSpecialTokens) { this.additionalSpecialTokens = additionalSpecialTokens; }

    public boolean isAddBosToken() { return addBosToken; }
    public String getBosToken() { return extractTokenContent(bosToken); }
    public String getEosToken() { return extractTokenContent(eosToken); }
    public String getPadToken() { return extractTokenContent(padToken); }
    public String getUnkToken() { return extractTokenContent(unkToken); }
    public String getChatTemplate() { return chatTemplate; }
    public int getModelMaxLength() { return modelMaxLength; }
    public String getTokenizerClass() { return tokenizerClass; }
    public Map<String, AddedToken> getAddedTokensDecoder() { return addedTokensDecoder; }
    public List<String> getAdditionalSpecialTokens() { return additionalSpecialTokens; }


    // Factory methods
    public static TransformersTokenizerConfig fromFile(String path) throws IOException {
        ObjectMapper mapper = new ObjectMapper();
        mapper.enable(com.fasterxml.jackson.databind.DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES);
        mapper.enable(com.fasterxml.jackson.core.JsonParser.Feature.ALLOW_UNQUOTED_FIELD_NAMES);
        mapper.enable(com.fasterxml.jackson.core.JsonParser.Feature.ALLOW_SINGLE_QUOTES);
        return mapper.readValue(new File(path), TransformersTokenizerConfig.class);
    }

    public static TransformersTokenizerConfig fromJson(String json) throws IOException {
        ObjectMapper mapper = new ObjectMapper();
        mapper.enable(com.fasterxml.jackson.databind.DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES);
        mapper.enable(com.fasterxml.jackson.core.JsonParser.Feature.ALLOW_UNQUOTED_FIELD_NAMES);
        mapper.enable(com.fasterxml.jackson.core.JsonParser.Feature.ALLOW_SINGLE_QUOTES);
        return mapper.readValue(json, TransformersTokenizerConfig.class);
    }

    public static TransformersTokenizerConfig fromInputStream(InputStream inputStream) throws IOException {
        ObjectMapper mapper = new ObjectMapper();
        mapper.enable(com.fasterxml.jackson.databind.DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES);
        mapper.enable(com.fasterxml.jackson.core.JsonParser.Feature.ALLOW_UNQUOTED_FIELD_NAMES);
        mapper.enable(com.fasterxml.jackson.core.JsonParser.Feature.ALLOW_SINGLE_QUOTES);
        return mapper.readValue(inputStream, TransformersTokenizerConfig.class);
    }

    // 处理复杂的 BOS token 对象
    @com.fasterxml.jackson.annotation.JsonSetter("bos_token")
    public void setBosTokenComplex(Object bosToken) {
        if (bosToken == null) {
            this.bosToken = null;
        } else if (bosToken instanceof String) {
            this.bosToken = (String) bosToken;
        } else if (bosToken instanceof Map) {
            @SuppressWarnings("unchecked")
            Map<String, Object> tokenMap = (Map<String, Object>) bosToken;
            this.bosToken = (String) tokenMap.get("content");
        }
    }

    // 处理复杂的 EOS token 对象
    @com.fasterxml.jackson.annotation.JsonSetter("eos_token")
    public void setEosTokenComplex(Object eosToken) {
        if (eosToken == null) {
            this.eosToken = null;
        } else if (eosToken instanceof String) {
            this.eosToken = (String) eosToken;
        } else if (eosToken instanceof Map) {
            @SuppressWarnings("unchecked")
            Map<String, Object> tokenMap = (Map<String, Object>) eosToken;
            this.eosToken = (String) tokenMap.get("content");
        }
    }


    // 创建 ChatTemplate 实例
    public ModelChatTemplate createChatTemplate() {
        return new ModelChatTemplate.Builder()
            .template(this.chatTemplate)
            .bosToken(this.bosToken != null ? this.bosToken : "")
            .eosToken(this.eosToken != null ? this.eosToken : "")
            .build();
    }
    public ModelChatTemplate createChatTemplate(boolean debug) {
        return new ModelChatTemplate.Builder()
            .template(this.chatTemplate)
            .bosToken(this.bosToken != null ? this.bosToken : "")
            .eosToken(this.eosToken != null ? this.eosToken : "")
            .debug(debug)
            .build();
    }
}
